import subprocess

import requests, time, os, re, json
from openpyxl import load_workbook
from requests_toolbelt import MultipartEncoder

def upgrade_dut(old_fw, new_fw, host, web_user, web_pwd, loop, logger, func):
    rsa_pwd = pass_k(host, web_pwd, logger)
    fw_upgrade(old_fw, new_fw, host, web_user, web_pwd, rsa_pwd, loop, logger, func)

def fw_upgrade(older_fw, newer_fw, host, user, pwd, rsa_pwd, loop, logger, func):
    result = {}
    # 创建报告目录
    report_path = os.getcwd() + '\\report\\' + time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))
    os.makedirs(report_path)
    # 加载模板文件
    wb = load_workbook(filename='template_fw.xlsx')
    ws = wb['fw']
    # 重启设备
    # reboot(host,user,pwd)
    # 循环执行固件升级测试
    for i in range(1, loop + 1):
        logger.info(f"第{i}次测试")
        func(f"第{i}次测试")
        result[i] = {}
        # 遍历要升级的固件文件
        for file in (older_fw, newer_fw):
            logger.info('开始升级，目标固件{0}'.format(file))
            func('开始升级，目标固件{0}'.format(file))
            # 获取固件版本号
            test_ver = re.findall('[0-9a-zA-Z]-(.+?)_apps.bin', file.split('/')[-1])[0]
            # 获取固件文件路径
            # fw_path = os.getcwd() + '\\FW\\' + file
            # 获取设备当前版本号
            ver = get_ver(host, user, rsa_pwd, logger)
            # 升级固件
            func("升级中，请耐心等待")
            rpc_log = upgrade(host, user, pwd, logger, file)
            # 等待设备恢复联网,每10s ping一次
            func('升级完成，设备重启中，等待恢复网络')
            time.sleep(60)
            while True:
                response = subprocess.run(['ping', '-n', '1', host], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                if b'TTL' in response.stdout:
                    logger.info(f"ping [{host}] 成功，网络恢复")
                    func(f"ping [{host}] 成功，网络恢复")
                    break
                else:
                    logger.info(f"ping [{host}] 失败,10s后继续检测网络")
                    func(f"ping [{host}] 失败,10s后继续检测网络")
                    time.sleep(10)
            func('设备重启成功，网络已恢复，验证升级结果')
            # 获取升级后的设备版本号
            ver1 = get_ver(host, user, rsa_pwd, logger)
            # 判断升级是否成功
            if ver1 == test_ver:
                result[i][test_ver] = 'Passed'
            else:
                result[i][test_ver] = 'Failed'
            # 保存升级日志
            result[i][test_ver + '_Log'] = rpc_log
            logger.info(f'测试固件版本:{test_ver}, 升级后系统版本:{ver1}, 测试结果:{result[i][test_ver]}')
            func(f'测试固件版本:{test_ver}, 升级后系统版本:{ver1}, 测试结果:{result[i][test_ver]}')
        # 判断本次循环的升级结果
        if 'Failed' not in list(result[i].values()) and len(list(result[i].values())) == 4:
            result[i]['Result'] = 'Passed'
        else:
            result[i]['Result'] = 'Failed'
        logger.info(f"第{i}次测试结束")
        func(f"第{i}次测试结束，测试结果:{result[i]['Result']}")
        # 保存测试结果到Excel文件
        judge = {}
        test_ver1 = re.findall('[0-9a-zA-Z]-(.+?)_apps.bin', older_fw.split('/')[-1])[0]
        test_ver2 = re.findall('[0-9a-zA-Z]-(.+?)_apps.bin', newer_fw.split('/')[-1])[0]
        logger.info(result)
        for get in result.items():
            temp = get[1]
            ws.cell(column=3, row=1).value = test_ver1
            ws.cell(column=4, row=1).value = test_ver1 + "_Log"
            ws.cell(column=5, row=1).value = test_ver2
            ws.cell(column=6, row=1).value = test_ver2 + "_Log"
            ws.cell(column=1, row=1 + get[0]).value = str(get[0])
            ws.cell(column=2, row=1 + get[0]).value = temp['Result']
            ws.cell(column=3, row=1 + get[0]).value = temp[test_ver1]
            ws.cell(column=4, row=1 + get[0]).value = temp[test_ver1 + '_Log']
            ws.cell(column=5, row=1 + get[0]).value = temp[test_ver2]
            ws.cell(column=6, row=1 + get[0]).value = temp[test_ver2 + '_Log']
            judge[str(get[0])] = temp['Result']
        wb.save('{}\\{}.xlsx'.format(report_path, 'fw_upgrade'))
    # # 判断整个测试过程的结果
    logger.info('upgarde result:{0}'.format(list(judge.values())))
    func ('测试结果:{0}'.format(list(judge.values())))
    # assert 'Failed' not in list(judge.values()) and len(list(judge.values())) == loop, 'upgrade test failed'

# 用于处理RSA加密密码
def pass_k(host, web_pwd, logger):
    from Crypto.PublicKey import RSA
    from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
    import base64
    se = requests.Session()
    url = f'http://{host}/cgi-bin/luci/guide/getPublicKey'
    y = se.post(url)
    try:
        import json
        key = json.loads(y.text)['rsaPublicKey']
        pub_key = RSA.importKey(str(key))
        cipher = PKCS1_cipher.new(pub_key)
        rsa_text = base64.b64encode(cipher.encrypt(bytes(web_pwd.encode("utf8"))))
        pass_s = rsa_text.decode('utf-8')
        logger.info(pass_s)
    except Exception as e:
        logger.info(f'{e},RSA not exist')
        pass_s = web_pwd
    return pass_s

# 定义请求头
HEADERS = {
    'Host': '192.168.10.1',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:82.0) Gecko/20100101 Firefox/82.0',
    'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
    'Accept-Language': 'en-US,en;q=0.5',
    'Accept-Encoding': 'gzip, deflate',
    'Content-Type': 'application/x-www-form-urlencoded',
    'Origin': 'http://192.168.10.1',
    'DNT': '1',
    'Referer': 'http://192.168.10.1/cgi-bin/luci/',
    'Upgrade-Insecure-Requests': '1'
}

def login(hostip, username, password, logger):
    # 创建会话
    se = requests.Session()
    # 定义认证信息
    auth = {"luci_acnt": username, "luci_pwd": password}
    # 更新请求头
    headers = HEADERS.update({'Host': hostip})
    headers = HEADERS.update({'Origin': f'http://{hostip}'})
    headers = HEADERS.update({'Referer': f'http://{hostip}/cgi-bin/luci/'})
    # 发送登录请求
    x = se.post(f'http://{hostip}/cgi-bin/luci/', headers=HEADERS, data=auth)
    # 获取登录Cookie
    # global SYSAUTH
    SYSAUTH = x.request.headers.get('Cookie', '')
    logger.info('GET login cookie - {0}'.format(SYSAUTH))
    return SYSAUTH

def get_ver(hostip, web_user, web_pwd, logger):
    # 登录并获取SYSAUTH
    SYSAUTH = login(hostip, web_user, web_pwd, logger)
    # 初始化数据字典
    data = {}
    # 创建requests会话
    se = requests.Session()
    # 设置请求头
    headers = HEADERS.update({'Host': hostip})
    headers = HEADERS.update({'Origin': f'http://{hostip}'})
    headers = HEADERS.update({'Referer': f'http://{hostip}/cgi-bin/luci/'})
    headers = HEADERS.update({'Referer': f'http://{hostip}/cgi-bin/luci/admin/starnet/more'})
    headers = HEADERS.update({'Cookie': SYSAUTH})
    try:
        url = f'http://{hostip}/cgi-bin/luci/admin/starnet/more/upgrade'
        # 发送POST请求
        y = se.post(url, headers=HEADERS, data=data)
        logger.info(y.status_code)
        # 使用正则表达式提取当前版本信息
        import re
        ex = '<span class="currentversion">(.+?)</span>'
        get = re.findall(ex, y.text)[0].replace('当前版本：', '')
        logger.info(get)
    except Exception as e:
        logger.info(e)
        get = ''
    return get

def get_ubus_session(host, web_user, web_pwd, logger):
    se = requests.Session()
    payload = {"jsonrpc": "2.0",
               "id": 4,
               "method": "call",
               "params": ["00000000000000000000000000000000", "session", "login",
                          {"username": web_user, "password": web_pwd}]
               }
    payload = json.dumps(payload)
    resp = se.post(f'http://{host}/ubus', data=payload, timeout=10)
    ubus_json = json.loads(resp.text)
    print(ubus_json)
    ubus_dict = ubus_json.get('result', '')[1]
    ubus_session = ubus_dict.get('ubus_rpc_session', '')
    # print(ubus_session)
    data = {"jsonrpc": "2.0",
            "id": 1,
            "method": "call",
            "params": [ubus_session, "system", "info", {}]
            }
    data = json.dumps(data)
    x = se.post(f'http://{host}/ubus', data=data, timeout=10)
    # test_info = json.loads(x.text)
    logger.info(x.text)
    test_info = x.text
    # if '[fx_sysupgrade] Memory too low. ' in "stdout":
    # logging.info('need reboot')
    # reboot(host,web_user,web_pwd)
    return ubus_session

# 定义upload函数，用于上传固件
def upload(host, fw, ubus_session, logger):
    from collections import OrderedDict
    # 打开固件文件
    fwFile = open(r'' + fw, 'rb')

    # 定义表单字段
    fields = OrderedDict([('sessionid', ubus_session),
                          ('filename', '/tmp/firmware.bin'),
                          ('filedata', (fw, fwFile, 'application/octet-stream'))])

    # 创建MultipartEncoder对象
    m = MultipartEncoder(
        fields=fields
    )
    # 定义请求头
    headers = {
        'Host': host,
        'Origin': 'http://%s/' % (host),
        'Referer': 'http://%s/cgi-bin/luci/admin/starnet/more/upgrade' % (host),
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36',
        'X-Requested-With': 'XMLHttpRequest',
        'Accept': '*/*',
        'Accept-Encoding': 'gzip, deflate',
        'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
        'Connection': 'keep-alive'
    }
    # 更新请求头的Content-Type
    headers.update({'Content-Type': m.content_type})

    se = requests.Session()
    x = se.post(f'http://{host}/cgi-bin/cgi-upload', headers=headers, data=m)
    logger.info(x.text)

# 定义jsonrpc_send函数，用于发送jsonrpc请求
def jsonrpc_send(host, params, logger):
    # 创建一个requests会话
    se = requests.Session()
    # 定义请求数据
    data = {"jsonrpc": "2.0",
            "id": 4,
            "method": "call",
            "params": params
            }
    # 将数据转换为JSON格式
    data = json.dumps(data)
    # 发送POST请求
    x = se.post(f'http://{host}/ubus', data=data, timeout=10)
    logger.info(x.text)
    return x.text

# 定义upgrade_fw函数，用于升级固件
def upgrade(host, web_user, web_pwd, logger, fw_path=r''):
    # 获取ubus会话
    ubus_session = get_ubus_session(host, web_user, web_pwd, logger)
    # 上传固件
    upload(host, fw_path, ubus_session, logger)
    # ubus_session = self.get_ubus_session()
    # params1 = [ubus_session, "system", "validate_firmware_image",{"path":"/tmp/firmware.bin"}]
    # self.jsonrpc_send(params1)
    time.sleep(5)
    # 获取ubus会话
    ubus_session = get_ubus_session(host, web_user, web_pwd, logger)
    try:
        # 定义请求参数
        params3 = [ubus_session, "file", "exec",
                   {"command": "/sbin/fx_sysupgrade", "params": ["/tmp/firmware.bin"], "env": {"HOME": "/root"}}]
        # 发送jsonrpc请求
        get = jsonrpc_send(host, params3, logger)
    except Exception as e:
        # 打印日志信息
        logger.info(f'dut reboot,{e},wait for 70s')
        # 设置get变量为'normal'
        get = 'normal'
    return get

